import itertools
import numpy as np

def load_reservoir_4sk_data():
    # Define the reservoir data

    y  = np.array([2, 3, 0, 0])
    a0 = np.array([5, 5, 5, 5])
    aN = np.array([5, 5, 5, 7])

    nb_mon = 12
    V = 4 * nb_mon

    u_max = np.array([3, 4, 4, 7])
    u_max = [element for element in u_max for _ in range(nb_mon)]
    u_max = np.array(u_max).reshape(4, nb_mon)

    s_max = np.array([10, 10, 10, 15])
    s_max = [element for element in s_max for _ in range(nb_mon)]
    s_max = np.array(s_max).reshape(4, nb_mon)

    s_min = np.zeros([4, nb_mon])
    u_min = np.zeros([4, nb_mon])

    b1 = [[1.1, 1.0, 1.0, 1.2, 1.8, 2.5, 2.2, 2.0, 1.8, 2.2, 1.8, 1.4]]
    b2 = [[1.4, 1.1, 1.0, 1.0, 1.2, 1.8, 2.5, 2.2, 2.0, 1.8, 2.2, 1.8]]
    b3 = [[1.0, 1.0, 1.2, 1.8, 2.5, 2.2, 2.0, 1.8, 2.2, 1.8, 1.4, 1.1]]
    b4 = [[1.0, 1.2, 1.8, 2.5, 2.2, 2.0, 1.8, 2.2, 1.8, 1.4, 1.1, 1.0]]
    b5 = [[1.6, 1.7, 1.8, 1.9, 2.0, 2.0, 2.0, 1.9, 1.8, 1.7, 1.6, 1.5]]

    b = np.vstack((b1, b2, b3, b4, b5))

    m_list = u_max.reshape(4, 12)[:, 1] + 1
    action_space = list(itertools.product(*(range(m) for m in m_list)))

    m_list = u_max.reshape(4, 12)[:, 1] + 1
    state_space = list(itertools.product(*(range(m) for m in m_list)))

    return y, a0, aN, u_max, u_min, s_max, s_min, b, nb_mon, action_space

if __name__ == '__main__':
    # Example usage:
    y, a0, aN, u_max, u_min, s_max, s_min, b, nb_mon, action_space = load_reservoir_4sk_data()
    print("Initial capacity a0:", a0)
    print("Final capacity aN:", aN)
    print("Maximum discharge u_max:", u_max)
    print("Minimum discharge u_min:", u_min)
    print("Minimum storage s_min:", s_min)
    print("Maximum storage s_max:", s_max)
    print("Benefit coefficients b:", b)